package org.ricks.net.action.context;

import org.ricks.net.action.context.head.HeadMeta;
import org.ricks.net.action.context.session.Session;
import org.ricks.net.action.rpc.RpcError;
import org.ricks.net.action.context.head.RpcHeadMeta;
import org.ricks.net.action.utils.DecoderException;
import org.ricks.net.action.parser.ProtocolParser;
import org.ricks.net.action.parser.TypeReference;
import org.ricks.net.action.utils.EncoderException;
import java.io.IOException;
import java.nio.ByteBuffer;
import java.util.concurrent.CompletableFuture;
import java.util.concurrent.TimeUnit;

/**
 * 上下文
 *
 * 无响应消息
 * rpcClient.send(CHAT_MSG_CMD, new ChatMessage("Hello"));
 *
 * 同步RPC调用
 * UserInfo user = rpcClient.invoke(GET_USER_CMD, userId, UserInfo.class).get(5, TimeUnit.SECONDS);
 *
 * 异步RPC调用（使用FlowContext）
 * flowContext.invoke(QUERY_DATA_CMD, params, ResultSet.class).thenAccept(this::processResults);
 */
public class FlowContext<T extends HeadMeta> {
    private final Session session;
    public final T headMeta;
    private final byte[] data;
    public final ProtocolParser parser;
    private final MessageBuilder messageBuilder;

    public FlowContext(Session session, T headMeta, byte[] data, ProtocolParser parser) {
        this.session = session;
        this.headMeta = headMeta;
        this.data = data;
        this.parser = parser;
        this.messageBuilder = new MessageBuilder(parser);
    }

    public int getCommand() {
        return headMeta.cmd();
    }

    public byte[] getData() {
        return data;
    }

    protected void doWrite(byte[] fullMessage) {
        session.writeBinary(fullMessage);
    }


    public void writeBody(Object body) {
        doWrite(messageBuilder.build(headMeta, body));
    }

    public void writeBody(HeadMeta headMeta, Object body) {
        doWrite(messageBuilder.build(headMeta, body));
    }


    public void writeBytes(byte[] data) {
        doWrite(messageBuilder.build(headMeta, data));
    }

    public void writeBytes(HeadMeta headMeta, byte[] data) {
        doWrite(messageBuilder.build(headMeta, data));
    }

    public <R> R deserialize(Class<R> clazz) {
        try {
            return parser.deserialize(data, clazz);
        } catch (IOException e) {
            throw new DecoderException("Deserialization failed", e);
        }
    }

    public <R> R deserialize(TypeReference<R> typeRef) {
        try {
            return parser.deserialize(data, typeRef.getType());
        } catch (IOException e) {
            throw new DecoderException("Deserialization failed", e);
        }
    }


    /**
     * 请求响应模式 - 请求端
     * @param cmd
     * @param body
     * @param resultType
     * @return
     * @param <R>
     */
    public <R> CompletableFuture<R> invoke(short cmd, Object body, Class<R> resultType) {
        return session.invoke(cmd, body, resultType);
    }

    // 新增带超时的调用
    public <R> CompletableFuture<R> invoke(short cmd, Object body, Class<R> resultType, long timeout, TimeUnit unit) {
        return session.invoke(cmd, body, resultType, timeout, unit);
    }

    /**
     * 请求响应 - 响应端
     * @param result
     */
    public void rpcResponse(Object result) {
        if (!(headMeta instanceof RpcHeadMeta reqHead)) {
            throw new IllegalStateException("Only RpcHeadMeta supports rpcResponse");
        }

        RpcHeadMeta respHead = new RpcHeadMeta(
                (short) reqHead.cmd(),
                true,
                false,
                reqHead.rpcRequestId()
        );
        doWrite(messageBuilder.build(respHead, result));
    }

    public void rpcError(int code, String message) {
        rpcResponse(new RpcError(code, message));
    }

    // 优化1: 静态内部类+线程安全
    private static class MessageBuilder {
        private final ProtocolParser parser;

        public MessageBuilder(ProtocolParser parser) {
            this.parser = parser;
        }

        public byte[] build(HeadMeta headMeta, Object body) {
            ByteBuffer headerBuffer = ByteBuffer.allocate(headMeta.headSize());
            byte[] header = generateHeader(headMeta);
            byte[] bodyData = serialize(body);
            return combine(header, bodyData);
        }

        public byte[] build(HeadMeta headMeta, byte[] body) {
            return combine(generateHeader(headMeta), body);
        }

        private byte[] generateHeader(HeadMeta headMeta) {
            ByteBuffer headerBuffer = ByteBuffer.allocate(headMeta.headSize());
            return headMeta.generateHeader(headerBuffer).array();
        }

        private byte[] serialize(Object data) {
            try {
                return parser.serialize(data);
            } catch (IOException e) {
                throw new EncoderException("Serialization failed", e);
            }
        }

        // 优化3: 使用System.arraycopy高效合并
        private byte[] combine(byte[] header, byte[] body) {
            byte[] full = new byte[header.length + body.length];
            System.arraycopy(header, 0, full, 0, header.length);
            System.arraycopy(body, 0, full, header.length, body.length);
            return full;
        }
    }
}
